'''
Author: caishuyang
Date: 2023-03-12 14:59:47
LastEditors: caishuyang
LastEditTime: 2023-03-13 17:07:56
Description: 
'''
import torch, os
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import transforms
from PIL import Image
import numpy as np
import torch.utils.data as Data  

class PipeDataset(Dataset):
	def __init__(self, DatasetFolderPath, ImgTransform, LabelTransform, ShowSample=False):
		self.DatasetFolderPath = DatasetFolderPath
		self.ImgTransform = ImgTransform
		self.LabelTransform = LabelTransform
		self.ShowSample = ShowSample
		self.SampleFolders = os.listdir(self.DatasetFolderPath)

	def __len__(self):
		return len(self.SampleFolders)

	def __getitem__(self, item):
		SampleFolderPath = os.path.join(self.DatasetFolderPath, self.SampleFolders[item])  # 样本文件夹路径
		FusionImgPath = os.path.join(SampleFolderPath, 'img.png')
		LabelImgPath = os.path.join(SampleFolderPath, 'label.png')
		FusionImg = Image.open(FusionImgPath)
		FusionImg = FusionImg.convert("RGB")
		LabelImg = Image.open(LabelImgPath)
		LabelImg=LabelImg.resize((160, 120),Image.ANTIALIAS)
		labimg=np.array(LabelImg)

		'''
		两类对象通道分离策略
		'''
		oupimg=np.array([labimg,1-labimg])

		FusionImg=FusionImg.resize((160, 120),Image.ANTIALIAS)
		
		FusionImg=self.ImgTransform(FusionImg)
		LabelImg=torch.Tensor(oupimg)

		return FusionImg,LabelImg

if __name__=="__main__":
	
	trans=transforms.ToTensor()
	datapath="Dataset\\Train"
	pipe=PipeDataset(datapath,trans,trans)
	trainiter=Data.DataLoader(pipe,32,shuffle=True,num_workers=0)
	for data in trainiter:
		features,target=data
		print(features)

	



